package demo.main;

import com.jucoin.infra.utils.LaneEnvUtil;
import demo.kafka.streams.utils.Constants;
import demo.kafka.streams.utils.KafkaUtil;
import demo.kafka.streams.utils.ThreadUtil;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.serialization.LongDeserializer;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.*;
import org.apache.kafka.streams.kstream.*;
import org.apache.kafka.streams.state.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Arrays;
import java.util.Collections;
import java.util.Properties;

/**
 * 单词计数应用
 * <ol>
 *     <li>生产者发送文本行到 {@value #TOPIC_TEXT_LINES} 主题</li>
 *     <li>消费者从 {@value #TOPIC_WORDS_WITH_COUNTS} 主题读取单词计数结果</li>
 *     <li>单词计数应用会将单词计数结果发送到 {@value #TOPIC_WORDS_WITH_COUNTS} 主题</li>
 * </ol>
 */
public class WordCountApp {

    private static final Logger logger = LoggerFactory.getLogger(WordCountApp.class);

    public static final String TOPIC_TEXT_LINES = "test.stan";
    public static final String TOPIC_WORDS_WITH_COUNTS = "test.stan.wordCount";
    public static final String STORE_NAME = "counts-store";
    private static KafkaStreams kafkaStreams;

    public static void main(final String[] args) throws Exception {
        KafkaUtil.startConsumer(TOPIC_WORDS_WITH_COUNTS, true);

        KafkaUtil.startProducer(TOPIC_TEXT_LINES, true);

        String appId = LaneEnvUtil.renameWithLane("word.count-app");
        String intTopic = LaneEnvUtil.renameWithLane(TOPIC_TEXT_LINES);
        String outTopic = LaneEnvUtil.renameWithLane(TOPIC_WORDS_WITH_COUNTS);
        Properties props = new Properties();
        // APPLICATION_ID_CONFIG 实际上就是消费组ID
        props.put(StreamsConfig.APPLICATION_ID_CONFIG, LaneEnvUtil.renameWithLane("word.count-app"));
        // 从最早的消息开始消费
        props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");

        props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, Constants.KAFKA_BOOTSTRAP_SERVERS);
        props.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.String().getClass());
        props.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.String().getClass());
        props.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000); // 1秒提交一次
        // 状态存储配置，挂载的持久化路径，可以是本地文件系统或分布式文件系统
//        props.put(StreamsConfig.STATE_DIR_CONFIG, "/mnt/rocksdb-data");

        StreamsBuilder streamsBuilder = new StreamsBuilder();

        KStream<String, String> textLines = streamsBuilder.stream(LaneEnvUtil.renameWithLane(TOPIC_TEXT_LINES));

        KStream<String, String> words = textLines.flatMapValues(textLine -> Arrays.asList(textLine.toLowerCase().split("\\W+")));

        KGroupedStream<String, String> groupedWords = words.groupBy((key, word) -> word);

        KTable<String, Long> wordCounts = groupedWords.count(Materialized.<String, Long, KeyValueStore<Bytes, byte[]>>as(STORE_NAME));

        wordCounts.toStream()
                .peek((word, count) -> logger.info("单词: " + word + ", 计数: " + count))
                .to(TOPIC_WORDS_WITH_COUNTS, Produced.with(Serdes.String(), Serdes.Long()));

        Topology topology = streamsBuilder.build();
        kafkaStreams = new KafkaStreams(topology, props);
        kafkaStreams.start();

        System.err.println("启动成功，应用ID: " + props.getProperty(StreamsConfig.APPLICATION_ID_CONFIG));


        printStreamStatus();
        printAllWordCount();
        printStreamStatus();

//        kafkaStreams.close();
        // 添加关闭钩子，优雅关闭
        Runtime.getRuntime().addShutdownHook(new Thread(kafkaStreams::close));

    }

    public static void printStreamStatus() {
        System.err.println("Kafka Streams 状态: " + kafkaStreams.state());
    }

    public static void printAllWordCount() {
        if (kafkaStreams == null) {
            return;
        }
        // 等待Streams完全启动再查询
        while (kafkaStreams.state() != KafkaStreams.State.RUNNING) {
            ThreadUtil.safeSleep(500);
        }
        try {
            QueryableStoreType<ReadOnlyKeyValueStore<String, Long>> readOnlyKeyValueStore = QueryableStoreTypes.keyValueStore();
            StoreQueryParameters<ReadOnlyKeyValueStore<String, Long>> storeQueryParameters =
                    StoreQueryParameters.fromNameAndType(STORE_NAME, readOnlyKeyValueStore);
            // 获取状态存储
            ReadOnlyKeyValueStore<String, Long> keyValueStore = kafkaStreams.store(storeQueryParameters);

            KeyValueIterator<String, Long> all = keyValueStore.all();
            while (all.hasNext()) {
                KeyValue<String, Long> next = all.next();
                System.err.println("单词: " + next.key + ", 计数: " + next.value);
            }
            all.close();

        } catch (Exception e) {
            System.err.println("查询失败: " + e.getMessage());
        }
    }


}
